import os
import subprocess
import sys
import json

# 添加当前目录到Python路径以便导入final_test
# 获取当前文件所在目录
current_dir = os.path.dirname(os.path.abspath(__file__))
# 将当前目录添加到Python路径
if current_dir not in sys.path:
    sys.path.append(current_dir)

# 导入final_test模块中的final_batch_pic函数
from final_test import final_batch_pic


def merge_frames_to_video(input_dir, output_video_path, fps=30):
    """
    使用ffmpeg将帧图像合并为视频文件
    
    参数:
    input_dir: 包含帧图像的目录
    output_video_path: 输出视频文件的路径
    fps: 帧率，默认为30
    """
    # 确保输入目录存在
    if not os.path.exists(input_dir):
        print(f"错误: 输入目录不存在: {input_dir}")
        return False
    
    # 检查输入目录中是否有图像文件
    frame_files = [f for f in os.listdir(input_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
    if not frame_files:
        print(f"错误: 输入目录中没有找到图像文件: {input_dir}")
        return False
    
    # 确保输出目录存在
    output_dir = os.path.dirname(output_video_path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"创建输出目录: {output_dir}")
    
    # 检查目录中实际的文件格式
    frame_files = [f for f in os.listdir(input_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
    if not frame_files:
        print(f"错误: 输入目录中没有找到图像文件: {input_dir}")
        return False
    
    # 判断文件命名格式
    sample_file = frame_files[0]
    if sample_file.startswith('2.0x-frame_') and sample_file.endswith('.png'):
        # Real-ESRGAN输出格式：2.0x-frame_000001.png
        input_pattern = '2.0x-frame_%06d.png'
    elif sample_file.startswith('frame_') and sample_file.endswith('.jpg'):
        # 标准格式：frame_000001.jpg
        input_pattern = 'frame_%06d.jpg'
    elif sample_file.startswith('frame_') and sample_file.endswith('.png'):
        # PNG格式：frame_000001.png
        input_pattern = 'frame_%06d.png'
    else:
        print(f"警告: 未识别的文件命名格式，使用第一个文件: {sample_file}")
        # 尝试通用格式
        input_pattern = sample_file.replace(sample_file.split('_')[-1], '%06d.' + sample_file.split('.')[-1])
    
    print(f"检测到文件格式: {input_pattern}")
    
    # 构建ffmpeg命令
    # -framerate: 设置输入帧率
    # -i: 输入文件模式（使用通配符匹配按顺序的帧文件）
    # -c:v: 视频编码器
    # -crf: 恒定速率因子（质量控制，值越小质量越高）
    # -pix_fmt: 像素格式
    ffmpeg_cmd = [
        'ffmpeg',
        '-framerate', str(fps),
        '-i', os.path.join(input_dir, input_pattern),
        '-c:v', 'libx264',
        '-crf', '23',
        '-pix_fmt', 'yuv420p',
        output_video_path, 
        '-y'
    ]
    
    print(f"开始合并帧为视频，帧率: {fps}")
    print(f"输入目录: {input_dir}")
    print(f"输出视频: {output_video_path}")
    print(f"执行命令: {' '.join(ffmpeg_cmd)}")
    
    try:
        # 执行ffmpeg命令
        result = subprocess.run(
            ffmpeg_cmd,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        print(f"视频合成完成! 输出文件: {output_video_path}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"合并帧为视频失败: {e}")
        print(f"错误输出: {e.stderr}")
        return False
    except Exception as e:
        print(f"发生未知错误: {e}")
        return False


def extract_frames(video_path, output_dir, fps=30):
    """
    使用ffmpeg从视频中提取帧
    
    参数:
    video_path: 视频文件路径
    output_dir: 输出帧文件夹路径
    fps: 帧率，默认为30
    """
    # 确保输出目录存在，如果不存在则创建
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"创建输出目录: {output_dir}")
    
    # 构建ffmpeg命令
    # -i: 输入文件
    # -r: 设置帧率
    # %06d: 6位数字序号，从000001开始
    ffmpeg_cmd = [
        'ffmpeg',
        '-i', video_path,
        '-r', str(fps),
        '-q:v', '2',  # 图像质量，数值越小质量越高
        os.path.join(output_dir, 'frame_%06d.jpg')
    ]
    
    print(f"开始提取视频帧，帧率: {fps}")
    print(f"视频文件: {video_path}")
    print(f"输出目录: {output_dir}")
    print(f"执行命令: {' '.join(ffmpeg_cmd)}")
    
    try:
        # 执行ffmpeg命令
        result = subprocess.run(
            ffmpeg_cmd,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        print("视频帧提取完成!")
        return True
    except subprocess.CalledProcessError as e:
        print(f"提取视频帧失败: {e}")
        print(f"错误输出: {e.stderr}")
        return False
    except Exception as e:
        print(f"发生未知错误: {e}")
        return False


if __name__ == '__main__':
    # 视频文件路径
    video_file = r"e:/py_test/final2x_media_handle/final_mdia/video_src/final2x_demo2.mp4"
    
    # 输出目录 - 在视频文件所在目录下创建frames子目录
    video_dir = os.path.dirname(video_file)
    output_directory = os.path.join(video_dir, 'frames')
    output_final2x = output_directory+'final2x'
    
    # # 检查视频文件是否存在
    # if not os.path.exists(video_file):
    #     print(f"错误: 视频文件不存在: {video_file}")
    #     sys.exit(1)
    
    # # 提取视频帧，帧率为25
    # success = extract_frames(video_file, output_directory, fps=25)
    success = 1 
    if success:
        # 统计提取的帧数
        try:
            frames = [f for f in os.listdir(output_directory) if f.endswith('.jpg')]
            print(f"共提取了 {len(frames)} 帧图像")
            # 获取帧文件的完整路径列表
            frame_paths = [os.path.join(output_directory, frame) for frame in frames]
            
            # 调用final_batch_pic函数处理提取的帧
            # print(f"开始调用final_batch_pic处理 {len(frame_paths)} 帧图像")
            # final_batch_pic(frame_paths, output_final2x)
            # 合并帧为视频
            # 修复路径拼接问题，使用os.path.join连接目录
            upscale_dir = os.path.join(video_dir, output_final2x+r'\outputs')
            output_video = os.path.join(video_dir, output_final2x+r'\final2x.mp4')
            merge_frames_to_video(upscale_dir, output_video, fps=5)
        except Exception as e:
            print(f"统计帧数时发生错误: {e}")
    
    print("程序执行完毕。")